#include<bits/stdc++.h>
#include<armadillo>
#include "Cell.h"
/*
 * compilar : make 1D_CN
 * ejecutar : ./run.sh 1D_CN
 *
 * Simulacion de la propagación del potencial de acción en una fibra usando
 * el método numérico Crank-Nicholson y el modelo celular auricular CRN98
 * de Courtemanche, Ramirez,Nattel
 */

using namespace std;
using namespace arma;
#define db double
#define PI 3.14159265

// Matrix A para el sistema AX=B.
void load_matriz_A(mat &A, int N, db r){
  for(int x=0; x<N; x++){
    // tiempo i=1, se cumple condicion de borde en i=0.
    A(x,x) = (1.0+2.0*r);
    if(x==0){
      A(x,x+1) = -r;
    }
    // Tiempo i=N, se cumple condicion de borde en i=N+1.
    else if(x==N-1){ //time N+1
      A(x,x-1) = -r;
    }
    // Diagonal para los otros tiempos.
    else{
      A(x,x-1) = -r;
      A(x,x+1) = -r;
    }
  }
}

// Copiar los voltages calculados por DF a cada celula.
void copy_voltage(vector<Cell> &cells, vec X){
  for(int i=1; i < X.size()+1; i++)
    cells[i].V = X(i-1);
}

// Imprime los voltages calculados por DF en un tiempo determinado.
void print_solutions(vec &X,db t){
  cout<<t;
  for(int i=0; i<X.n_rows; i++)
    cout<<"     "<<X(i);
  cout<<endl;
}

int main(){
  db D;           // coeficiente de difusion
  int N;          // Numero de celulas a simular
  db nrepeat;     // numero de ciclos
  db tbegin;      // tiempo de inicio del primer estímulo
  db BCL;         // frecuencia de excitacion en mse
  db CI;          // intervalo de acoplamiento para el ultimo
  db dt;          // paso de tiempo
  db dtstim;      // duracion del estimulo
  db CurrStim;    // corriente de estimulo
  int nstp_prn;   // frecuencia con la que se imprimen los resultados
  db nstp;
  db Iion;
  db Jion;          // current density
  db beta;
  db cell_to_stim;  // numero de la celula a estimular
  db deltaX;
  db deltaX2;
  db t=0.0;
  db tend;
  bool flag_stm=1;
  db Istim=0.0;
  db cont_repeat = 0;

  // parametros del sistema
  N = 20;
  nrepeat = 1; // number of beats
  tbegin = 50;  // 100; //50
  BCL =  600;
  CI = 0;
  dt = 0.02;                  //[ms]
  deltaX = 0.025;             //[cm]
  deltaX2 = deltaX*deltaX;
  dtstim = 2;
  CurrStim = -8000;
  nstp_prn = 10;
  cell_to_stim = 1;

  tend = tbegin+dtstim;
  nstp = (tbegin+BCL*nrepeat+CI)/dt;

  vector<Cell> cells(N+2);

  db areaT = cells[0].pi*pow(cells[0].a,2);    // Capacitive membrane area
  db aCm = cells[0].Cap / areaT;               // Capacitance per unit area pF/cm^2

  D = cells[0].a / (2.0*cells[0].Ri*aCm*1e-9); //D = 0.00217147 [cm^2/ms]

  mat A = mat(N,N);           // A
  vec B = vec(N);             // B
  vec X = vec(N);             // X from AX=B;

  db r = (D*dt)/(2.0*deltaX2);
  load_matriz_A(A, N, r);

  vec preV= vec(N+2);
  preV.fill(-81.2);

//var for printing only the last ncharts beats
  int ncharts = 1;
  int time_to_print = nstp- ((ncharts*BCL+tbegin)/dt);

  for(int k=0; k <nstp+2; k++,t+=dt){       //iteracion sobre el tiempo
    if(t>=tbegin && t<=tend){
      flag_stm = 0;
    }
    else{
      if(flag_stm==0){
        if(cont_repeat < nrepeat){
          tbegin=tbegin+BCL; //se establece el tiempo del próximo estimulo
        }
        else if(cont_repeat == nrepeat) tbegin=tbegin+CI;
        cont_repeat++;
        tend=tbegin+dtstim;
        flag_stm = 1;
      }
    }

    // Metodo Crank- Nicholson
    for(int x=1; x<N+1; x++){
      // La corriente de estimulo es aplicada en un periodo de tiempo.
      if(!flag_stm && x == cell_to_stim){
        Istim = CurrStim;
      }else{
        Istim = 0.0;
      }

      Iion = cells[x].getItot(dt);
      if(x==1){
        beta = (1.0-2.0*r)*preV(x) + r*preV(x+1) +2*r*preV(0);
      }
      else if(x==N){
        beta = r*preV(x-1) + (1.0-2.0*r)*preV(x) + 2*r*preV(N+1);
      }
      else{
        beta= r*preV(x-1) + (1.0-2.0*r)*preV(x) + r*preV(x+1);
      }
      Jion = (Iion+Istim)/areaT;
      B(x-1) = beta - (dt/aCm)*(Jion);
    }
    // Armadillo: Solver for AX=B
    X = solve(A,B);
    preV.subvec(1,N) = X;

    //Copiar voltajes a las celulas.
    copy_voltage(cells,X);
    if(k%nstp_prn==0 && k>time_to_print)   //use this for plot last beat
      //if(k%nstp_prn==0)                  //use this for plot all beats
      print_solutions(X,t);
  }
  return 0;
}
